#This files takes functions from MCES..._functions and adapts to Mixed Logit case 
#when relevant.
#1) settings.mlog, 2) eqbm.mlog 3) monte.func.mlog and monte.rep.mlog that use 1) and 2)

####################################
#Equilibrium p and s on each market#
####################################
eqbm.mlog <- function(dat,market,v=1,tar.change=0,verbose=FALSE) { 
  #
  source("eqbm_MLOG_subfunctions.R",local=TRUE)
  #
  ###########################
  #start of eqbm computation#  
  df.m <- data.frame(subset(dat$models,n==market),tar.change=tar.change)
  df.m <- within(df.m,{tar.change[i==n]<-0})
  i <- df.m$i
  if(U0==0) xvars <- paste0("x",1:n.x) else xvars <- paste0("x",0:n.x)
  x <- as.matrix(df.m[,xvars]) # if df.m is only a data.frame
  colnames(x) <- xvars  # no effect for n.x>1, but names it x1 instead of just x
  xi <- df.m$xi
  mc <- df.m$mc * (df.m$tau + df.m$tar.change*df.m$tariff.change) # apply trade costs and name it mc, for consistency with subordinate functions
  #
  df <- subset(dat$households,n==market)
  if(U0==0) bvec <- paste0("b",1:n.x,".h") else bvec <- paste0("b",0:n.x,".h")
  b.h <- as.matrix(df[,bvec]) 
  colnames(b.h) <- bvec  # no effect for n.x>1, but names it b1.h instead of just b.h
  alpha.h <- df$alpha.h
  y.h <- df$y.h
  Y <- sum(df$y.h)
  co.own <- co.own.mat
  # FPI to get the MCES predict prices and market shares
  FPIresults <- fpiter(mc,fixptfn=price.update,nu=settings.var$speed[v],control=list(trace=verbose,maxiter=MAXIT))
  n.loops  <- FPIresults$fpevals
  p <- FPIresults$par
  s <- colMeans(prob.mat(p))
  pelas.m <- ownelas(prm=prob.mat(p),share=s,price=p) # model level own price elasticities
  superelas.m <- superelas(prm=prob.mat(p),share=s,price=p) # model level own super elasticity
  ptr.sp.m <- (pelas.m)/(pelas.m + 1 - superelas.m)
  pte.sp.m <- (pelas.m + 1)/(pelas.m + 1 - superelas.m)
  #average income per model and ISE
  avg.inc.m <- avg.inc(prob.mat(p))   
  avg.inc.results <- lm(log(avg.inc.m)~log(p))
  ISE <- as.numeric(avg.inc.results$coefficients[2])
  return(data.frame(m=df.m$m,f=df.m$f,i,n=market,x,xi,mc=df.m$mc,cit=mc,s,p,pelas.m,ptr.sp.m,pte.sp.m,avg.inc.m,ISE,n.loops=n.loops,Y=Y))
  }

#Putting all markets together
DGP.endog.mlog <- function(parallel=TRUE,v=1,tar.change=0,data.exog){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.countries, .combine = 'rbind')  %dorng%  eqbm.mlog(dat=data.exog,i,v,tar.change) else foreach(i=1:n.countries, .combine = 'rbind')  %do%  eqbm.mlog(dat=data.exog,i,v,tar.change) 
}

##########################
######## MC part #########

#Monte Carlo repetitions MM - MC
monte.func.mlog.mm <- function(repi,settings,v=1,market=1) { # monte carlo stage
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  #pre
  eqbm.pre.mm  <- as.data.table(DGP.endog.mlog(parallel=F,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results
  #post:
  eqbm.post <- eqbm.mlog(dat=data.exog.mm,market=market,v=v,tar.change=1)
  eqbm.pre.with.ces <- eqbm.pre.mm.with.ces[n==market,] #keep only one destination
  eqbm.pre <- eqbm.pre.mm[n==market,] #keep only one destination
  #AFTER THIS, nothing in the code is changed compared to 2ctry, since we have only one destination left.
  eqbm.post.with.ces <- as.data.table(CF.approx(eqbm.pre.with.ces,eqbm.post)) #add ces approximation results
  return(data.frame(repi=repi,v=v,CF.compare(eqbm.pre.with.ces,eqbm.post.with.ces)))
}

#Doing multiple monte carlo reps for a single market
monte.rep.mlog.mm <- function(parallel=TRUE,settings,v=1,market=1){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.reps, .combine = 'rbind')  %dorng%  monte.func.mlog.mm(i,settings,v,market) else foreach(i=1:n.reps, .combine = 'rbind')  %do%  monte.func.mlog.mm(i,settings,v,market) 
}

#OLY version 
monte.func.cf.mlog.oly.mm <- function(repi,settings,v=1,market=1) { # monte carlo stage
data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
#pre:
eqbm.pre.mm  <- as.data.table(DGP.endog.mlog(parallel=F,v=v,tar.change=0,data.exog = data.exog.mm))
eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results
eqbm.pre.mm.with.ces[, s.f := sum(s), by=c("f","n")]
eta.ces.oly.init <<- 3
FPIresults <- fpiter(eta.ces.oly.init,fixptfn=ces.oly.mm.update,pre=eqbm.pre.mm.with.ces,nu=1) #estimate eta.ces.oly
eta.ces.oly <- FPIresults$par 
eta.ces <- eqbm.pre.mm.with.ces$eta[1] 
eta.oly.choice <<- eta.ces.oly
eqbm.pre.with.ces.oly <- as.data.table(estapprox.oly.mm(eqbm.pre.mm.with.ces,approx.results=ols.results,eta.choice =  eta.oly.choice)) #add ces.oly approximation results
#post:
eqbm.post <- eqbm.mlog(dat=data.exog.mm,market=market,v=v,tar.change=1) #compute eqbm on one market after removing tariffs
eqbm.pre.with.ces <- eqbm.pre.mm.with.ces[n==market,] #keep only one destination
eqbm.pre.with.ces.oly <- eqbm.pre.with.ces.oly[n==market,] #keep only one destination
eqbm.pre <- eqbm.pre.mm[n==market,] #keep only one destination
#AFTER THIS, nothing in the code is changed compared to 2ctry, since we have only one destination left.
eqbm.post.with.ces <- as.data.table(CF.approx(eqbm.pre.with.ces,eqbm.post)) #add ces approximation results
eqbm.post.with.ces.oly <- as.data.table(CF.approx.oly(pre=eqbm.pre.with.ces,post=eqbm.post.with.ces,settings=settings,verbose = FALSE))
return(data.frame(repi=repi,v=v,CF.compare.oly(eqbm.pre.with.ces.oly,eqbm.post.with.ces.oly)))
}

#Doing multiple monte carlo reps for a single market
monte.rep.mlog.oly.mm <- function(parallel=TRUE,settings,v=1,market=1){ #get equilibrium market shares and prices for each market and append them
  if(parallel==T)  foreach(i=1:n.reps, .combine = 'rbind')  %dorng%  monte.func.cf.mlog.oly.mm(i,settings,v,market) else foreach(i=1:n.reps, .combine = 'rbind')  %do%  monte.func.cf.mlog.oly.mm(i,settings,v,market) 
}



#Graph data
graph.model.level.func.mm.mlog <- function(settings,v=1,market=1) { 
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  #pre
  eqbm.pre.mm  <- as.data.table(DGP.endog.mlog(parallel=F,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm.with.ces <- as.data.table(estapprox.mm(eqbm.pre.mm)) #add ces approximation results
  #post:
  eqbm.post.mm <- as.data.table(DGP.endog.mlog(parallel=F,v=v,tar.change=1,data.exog = data.exog.mm))
  pre.with.ces <- eqbm.pre.mm.with.ces[n==market,] #keep only one destination
  pre <- eqbm.pre.mm[n==market,] #keep only one destination
  post <- eqbm.post.mm[n==market,] #keep only one destination
  #AFTER THIS, nothing in the code is changed compared to 2ctry, since we have only one destination left.
  post.with.ces <- as.data.table(CF.approx(pre.with.ces,post)) #add ces approximation results
  #compare (after rescaling inside good)
  #
  pre$s <- pre$s/sum(pre$s)  
  pre.with.ces$s <- pre.with.ces$s/sum(pre.with.ces$s)  
  pre.with.ces$s.ces <- pre.with.ces$s.ces/sum(pre.with.ces$s.ces)  
  post.with.ces$s <- post.with.ces$s/sum(post.with.ces$s)  
  post.with.ces$s.ces <- post.with.ces$s.ces/sum(post.with.ces$s.ces)  
  #
  s.post <- post.with.ces$s
  s.ces.post <- post.with.ces$s.ces
  s.delta <-s.post-pre$s
  s.ces.delta <-s.ces.post-pre$s
  #compute PT
  tau.hat <- post$cit/pre$cit
  ptr <- ((post$p-pre$p)/(post$cit-pre$cit)) #  pass thru elasticity  a la Ind Org (IO)
  pte <- ((post$p-pre$p)/(post$cit-pre$cit))*(pre$cit/pre$p) #  pass thru elasticity  a la Int'l Macro (IM)
  ptr.theory <- pre$ptr.sp.m
  pte.theory <- pre$pte.sp.m
  rm(pre,post)
  return(data.table(v=v,pre.with.ces,s.post,s.ces.post,s.delta,s.ces.delta,ptr,pte,ptr.theory,pte.theory,tau.hat))
}


graph.model.level.func.mm.oly.mlog <- function(settings,v=1,market=1) { 
  data.exog.mm <- DGP.exog(settings=settings,v=v,grid=subset(tau.grid.allv,vtau==v,select=-vtau)) #gen exog data
  #pre:
  eqbm.pre.mm  <- as.data.table(DGP.endog.mlog(parallel=T,v=v,tar.change=0,data.exog = data.exog.mm))
  eqbm.pre.mm[, s.f := sum(s), by=c("f","n")]
  eta.ces.oly.init <<- 3
  FPIresults <- fpiter(eta.ces.oly.init,fixptfn=ces.oly.mm.update,pre=eqbm.pre.mm,nu=1) #estimate eta.ces.oly
  eta.ces.oly <- FPIresults$par 
  eta.oly.choice <<- eta.ces.oly
  eqbm.pre.oly.mm <- as.data.table(estapprox.oly.mm(eqbm.pre.mm,approx.results=ols.results,eta.choice =  eta.oly.choice)) #add ces.oly approximation results
  #post:
  post <- eqbm.mlog(dat=data.exog.mm,market=market,v=v,tar.change=1) #compute eqbm on one market after removing tariffs
  pre <- eqbm.pre.mm[n==market,] #keep only one destination
  pre.with.ces.oly <- eqbm.pre.oly.mm[n==market,] #keep only one destination
  post.with.ces.oly <- as.data.table(CF.approx.oly(pre=pre,post=post,settings=settings,verbose = FALSE))
  #compare (after rescaling inside good)
  #
  pre$s <- pre$s/sum(pre$s)  
  pre.with.ces.oly$s <- pre.with.ces.oly$s/sum(pre.with.ces.oly$s)  
  pre.with.ces.oly$s.ces.oly <- pre.with.ces.oly$s.ces.oly/sum(pre.with.ces.oly$s.ces.oly)  
  post.with.ces.oly$s <- post.with.ces.oly$s/sum(post.with.ces.oly$s)  
  post.with.ces.oly$s.ces.oly <- post.with.ces.oly$s.ces.oly/sum(post.with.ces.oly$s.ces.oly)  
  #
  s.post <- post.with.ces.oly$s
  s.ces.post <- post.with.ces.oly$s.ces.oly
  s.delta <-s.post-pre$s
  s.ces.delta <-s.ces.post-pre$s
  p.post <- post.with.ces.oly$p
  cit.post <- post.with.ces.oly$cit
  #compute PT
  pte <- ((post$p-pre$p)/(post$cit-pre$cit))*(pre$cit/pre$p) #  pass thru elasticity  a la Int'l Macro (IM)
  rm(pre,post)
  return(data.table(v=v,pre.with.ces.oly,s.post,s.ces.post,s.delta,s.ces.delta,pte,p.post,cit.post))
}


